#include "sgSuperLU.h"
//#include "superlu/slu_ddefs.h"
#include "slu_ddefs.h"
// #include "Eigen/SuperLUSupport"

using namespace Algebra;

/*
 * Reform from the snippet in SuperLU/EXAMPLE/dlinsolx1.cpp.
 *
 * This example illustrates how to use DGSSVX to solve systems with the same
 * A but different right-hand side.
 * In this case, we factorize A only once in the first call to DGSSVX,
 * and reuse the following data structures in the subsequent call to DGSSVX:
 *     perm_c, perm_r, R, C, L, U.
 */

namespace Algebra
{
    struct sgSuperLUPrivate
    {
        SuperMatrix       A, L, U;
        superlu_options_t options;
        // SuperLUStat_t stat;

        char     equed[1];
        yes_no_t equil = YES;
        trans_t  trans = NOTRANS;

        GlobalLU_t Glu;           /* facilitate multiple factorizations with
                                     SamePattern_SameRowPerm                  */
        int_t    *perm_c = nullptr; /* column permutation vector */
        int_t    *perm_r = nullptr; /* row permutations from partial pivoting */
        int_t    *etree  = nullptr;
        void   *work   = nullptr;
        int_t   info   = 0;
        int_t   lwork  = 0;
        double *R      = nullptr;
        double *C      = nullptr;
        double *ferr   = nullptr;
        double *berr   = nullptr;
        double  u;
        double  rpg;
        double  rcond;
    };
}  // namespace Algebra

//=============================================================================
// sgSuperLU

sgSuperLU::sgSuperLU ()
{
    d = new sgSuperLUPrivate ();

    init ();
}

sgSuperLU::~sgSuperLU ()
{
    clearFactors ();
    if (d)
        delete d;
}

void sgSuperLU::compute (const SpMat_t &matrix)
{
    // the solver can be re-used to another matrix equations
    // Thus, clear the old LU factorization.
    clearFactors ();

    /* Set the default values for options argument:
    options.Fact = DOFACT;
        options.Equil = YES;
        options.ColPerm = COLAMD;
    options.DiagPivotThresh = 1.0;
        options.Trans = NOTRANS;
        options.IterRefine = NOREFINE;
        options.SymmetricMode = NO;
        options.PivotGrowth = NO;
        options.ConditionNumber = NO;
        options.PrintStat = YES;
    */
    set_default_options (&d->options);
    d->options.Equil           = d->equil;
    d->options.DiagPivotThresh = d->u;
    d->options.Trans           = d->trans;

    if (d->lwork > 0)
    {
        d->work = SUPERLU_MALLOC (d->lwork);
        if (!d->work)
        {
            ABORT ("DLINSOLX: cannot allocate work[]");
        }
    }

    auto &A    = const_cast<SpMat_t &> (matrix);
    int_t   m    = A.rows ();      // the number of rows of the matrix
    int_t   n    = A.cols ();      // the number of columns of the matrix
    int_t   nnz  = A.nonZeros ();  // the number of non zero coefficients
    int_t   nrhs = 1;

    // The program first initializes the three arrays, a[], asub[] and xa[],
    // which store the nonzero coefficients of matrix A, their row indices,
    // and the indices indicating the beginning of each column in the coefficient
    // and row index arrays. This storage format is called compressed column format.

    double *a    = A.valuePtr ();       // a non-const pointer to the array of values.
    int    *asub = A.innerIndexPtr ();  // a non-const pointer to the array of inner indices.
    int    *xa   = A.outerIndexPtr ();  // a const pointer to the array of the starting positions
                                        // of the inner vectors.

    /* Create matrix A in the format expected by SuperLU. */
    dCreate_CompCol_Matrix (&d->A, m, n, nnz, a, asub, xa, SLU_NC, SLU_D, SLU_GE);

    NCformat *Astore = ( NCformat * ) d->A.Store;
#ifdef SG_VERBOSE
    printf ("Dimension %dx%d; # nonzeros %d\n", ( int ) d->A.nrow, ( int ) d->A.ncol, ( int ) Astore->nnz);
#endif

    double *rhsb = nullptr;
    double *rhsx = nullptr;
    if (!(rhsb = doubleMalloc (m * nrhs)))
        ABORT ("Malloc fails for rhsb[].");
    if (!(rhsx = doubleMalloc (m * nrhs)))
        ABORT ("Malloc fails for rhsx[].");

    SuperMatrix B, X;
    dCreate_Dense_Matrix (&B, m, nrhs, rhsb, m, SLU_DN, SLU_D, SLU_GE);
    dCreate_Dense_Matrix (&X, m, nrhs, rhsx, m, SLU_DN, SLU_D, SLU_GE);

    if (!(d->etree = intMalloc (n)))
        ABORT ("Malloc fails for etree[].");
    if (!(d->perm_r = intMalloc (m)))
        ABORT ("Malloc fails for perm_r[].");
    if (!(d->perm_c = intMalloc (n)))
        ABORT ("Malloc fails for perm_c[].");
    if (!(d->R = ( double * ) SUPERLU_MALLOC (d->A.nrow * sizeof (double))))
        ABORT ("SUPERLU_MALLOC fails for R[].");
    if (!(d->C = ( double * ) SUPERLU_MALLOC (d->A.ncol * sizeof (double))))
        ABORT ("SUPERLU_MALLOC fails for C[].");
    if (!(d->ferr = ( double * ) SUPERLU_MALLOC (nrhs * sizeof (double))))
        ABORT ("SUPERLU_MALLOC fails for ferr[].");
    if (!(d->berr = ( double * ) SUPERLU_MALLOC (nrhs * sizeof (double))))
        ABORT ("SUPERLU_MALLOC fails for berr[].");

    /* Initialize the statistics variables. */
    SuperLUStat_t stat;
    StatInit (&stat);

    mem_usage_t mem_usage;

    /* ONLY PERFORM THE LU DECOMPOSITION */
    B.ncol = 0; /* Indicate not to solve the system */
    dgssvx (&d->options,
            &d->A,
            d->perm_c,
            d->perm_r,
            d->etree,
            d->equed,
            d->R,
            d->C,
            &d->L,
            &d->U,
            d->work,
            d->lwork,
            &B,
            &X,
            &d->rpg,
            &d->rcond,
            d->ferr,
            d->berr,
            &d->Glu,
            &mem_usage,
            &stat,
            &d->info);
#ifdef SG_VERBOSE
    printf ("LU factorization: dgssvx() returns info %lld\n", ( long long ) d->info);
#endif

    if (d->info == 0 || d->info == n + 1)
    {
#ifdef SG_VERBOSE
        if (d->options.PivotGrowth)
            printf ("Recip. pivot growth = %e\n", d->rpg);
        if (d->options.ConditionNumber)
            printf ("Recip. condition number = %e\n", d->rcond);
        SCformat *Lstore = ( SCformat * ) d->L.Store;
        NCformat *Ustore = ( NCformat * ) d->U.Store;
        printf ("No of nonzeros in factor L = %lld\n", ( long long ) Lstore->nnz);
        printf ("No of nonzeros in factor U = %lld\n", ( long long ) Ustore->nnz);
        printf ("No of nonzeros in L+U = %lld\n", ( long long ) Lstore->nnz + Ustore->nnz - n);
        printf ("FILL ratio = %.1f\n", ( float ) (Lstore->nnz + Ustore->nnz - n) / nnz);

        printf ("L\\U MB %.3f\ttotal MB needed %.3f\n", mem_usage.for_lu / 1e6, mem_usage.total_needed / 1e6);
        fflush (stdout);
#endif
    }
    else if (d->info > 0 && d->lwork == -1)
    {
        printf ("** Estimated memory: %lld bytes\n", ( long long ) d->info - n);
    }

    if (d->options.PrintStat)
        StatPrint (&stat);
    StatFree (&stat);

    Destroy_SuperMatrix_Store (&B);
    Destroy_SuperMatrix_Store (&X);

    // Now, the LU factorization has been finished.
    m_factorizationIsOk = true;
}

Vec_t sgSuperLU::solve (const Vec_t &b) const
{
    assert (m_factorizationIsOk && "The decomposition is not in a valid state for solving, you must first call compute()");

    /* Create right-hand side matrix B. */
    int m    = d->A.nrow;
    int n    = d->A.ncol;
    int nrhs = 1;

    double *rhsb = nullptr;
    double *rhsx = nullptr;
    // double *xact = nullptr;
    if (!(rhsb = doubleMalloc (m * nrhs)))
        ABORT ("Malloc fails for rhsb[].");
    if (!(rhsx = doubleMalloc (m * nrhs)))
        ABORT ("Malloc fails for rhsx[].");

    // fill the rhsb and rhsx
    for (int i = 0; i < m; ++i)
    {
        rhsb[i] = b[i];
        rhsx[i] = 0.0;
    }

    SuperMatrix B, X;
    dCreate_Dense_Matrix (&B, m, nrhs, rhsb, m, SLU_DN, SLU_D, SLU_GE);
    dCreate_Dense_Matrix (&X, m, nrhs, rhsx, m, SLU_DN, SLU_D, SLU_GE);

    /* ------------------------------------------------------------
        NOW WE SOLVE THE LINEAR SYSTEM USING THE FACTORED FORM OF A.
        ------------------------------------------------------------*/
    d->options.Fact = FACTORED; /* Indicate the factored form of A is supplied. */
    // B.ncol = nrhs;  /* Set the number of right-hand side */

    /* Initialize the statistics variables. */
    SuperLUStat_t stat;
    StatInit (&stat);

    mem_usage_t mem_usage;

    dgssvx (&d->options,
            &d->A,
            d->perm_c,
            d->perm_r,
            d->etree,
            d->equed,
            d->R,
            d->C,
            &d->L,
            &d->U,
            d->work,
            d->lwork,
            &B,
            &X,
            &d->rpg,
            &d->rcond,
            d->ferr,
            d->berr,
            &d->Glu,
            &mem_usage,
            &stat,
            &d->info);

#ifdef SG_VERBOSE
    printf ("Triangular solve: dgssvx() returns info %lld\n", ( long long ) d->info);
#endif

    // fetch the final solution
    Vec_t x (m);

    if (d->info == 0 || d->info == n + 1)
    {
        /* This is how you could access the solution matrix. */
        double *sol = ( double * ) (( DNformat * ) X.Store)->nzval;
        ( void ) sol;  // suppress unused variable warning

#ifdef SG_VERBOSE
        if (d->options.IterRefine)
        {
            printf ("Iterative Refinement:\n");
            printf ("%8s%8s%16s%16s\n", "rhs", "Steps", "FERR", "BERR");
            for (int i = 0; i < nrhs; ++i)
                printf ("%8d%8d%16e%16e\n", i + 1, stat.RefineSteps, d->ferr[i], d->berr[i]);
        }
        fflush (stdout);
#endif

        // fetch the final solution
        for (int i = 0; i < m; ++i)
        {
            x[i] = sol[i];
        }
    }
    else if (d->info > 0 && d->lwork == -1)
    {
        printf ("** Estimated memory: %lld bytes\n", ( long long ) d->info - n);
    }

    if (d->options.PrintStat)
        StatPrint (&stat);
    StatFree (&stat);

    SUPERLU_FREE (rhsb);
    SUPERLU_FREE (rhsx);
    Destroy_SuperMatrix_Store (&B);
    Destroy_SuperMatrix_Store (&X);

    return x;
}

Info_t sgSuperLU::info () const
{
    return ( Info_t ) d->info;
}

void sgSuperLU::init ()
{
    m_factorizationIsOk = false;

    d->info    = Info_t::InvalidInput;
    d->L.Store = nullptr;
    d->U.Store = nullptr;
}

void sgSuperLU::clearFactors ()
{
    if (d == nullptr)
    {
        return;
    }

    SUPERLU_FREE (d->etree);
    SUPERLU_FREE (d->perm_r);
    SUPERLU_FREE (d->perm_c);
    SUPERLU_FREE (d->R);
    SUPERLU_FREE (d->C);
    SUPERLU_FREE (d->ferr);
    SUPERLU_FREE (d->berr);
    // Destroy_CompCol_Matrix(&d->A);

    if (d->lwork == 0)
    {
        if (d->L.Store)
            Destroy_SuperNode_Matrix (&d->L);
        if (d->U.Store)
            Destroy_CompCol_Matrix (&d->U);

        d->L.Store = 0;
        d->U.Store = 0;

        memset (&d->L, 0, sizeof d->L);
        memset (&d->U, 0, sizeof d->U);
    }
    else if (d->lwork > 0)
    {
        SUPERLU_FREE (d->work);
    }

    // flag that LU factorization has been not performed.
    m_factorizationIsOk = false;
}